#!/usr/bin/env python3
from lxml import etree
import os
import shutil
import argparse

DEBUG_MODE=True

def remove_module(dir, module_list, pom_file="pom.xml"):        
    pom_path = os.path.join(dir, pom_file)
    if DEBUG_MODE:
        print(f"模块列表: {module_list}")  # 打印模块列表
        print(f"POM 文件路径: {pom_path}")  # 打印 POM 文件路径

    # 解析 XML 文件
    tree = etree.parse(pom_path)
    root = tree.getroot()

    # 定义命名空间
    namespace = {'ns': 'http://maven.apache.org/POM/4.0.0'}

    # 查找 <modules> 元素
    modules = root.find('ns:modules', namespaces=namespace)
    if modules is not None:
        for module in modules.findall('ns:module', namespaces=namespace):
            if DEBUG_MODE:
                print(f"检查模块: {module.text}")
            if module.text in module_list:
                if DEBUG_MODE:
                    print(f"删除模块: {module.text}")
                modules.remove(module)
                module_directory = os.path.join(dir, module.text)
                if os.path.exists(module_directory) and os.path.isdir(module_directory):
                    if DEBUG_MODE:
                        print(f"删除目录: {module_directory}")
                    shutil.rmtree(module_directory)

    # 写回文件，保留原始格式
    with open(pom_path, 'wb') as f:
        f.write(etree.tostring(tree, pretty_print=True, encoding='utf-8', xml_declaration=True))

    if DEBUG_MODE:
        print(f"POM 文件已更新: {pom_path}")

def add_module(dir, module_list, pom_file="pom.xml"):
    add_list = module_list    
    pom_path = os.path.join(dir, pom_file)
    if DEBUG_MODE:
        print(f"模块列表: {module_list}")  # 打印模块列表
        print(f"POM 文件路径: {pom_path}")  # 打印 POM 文件路径

    # 解析 XML 文件
    tree = etree.parse(pom_path)
    root = tree.getroot()

    # 定义命名空间
    namespace = {'ns': 'http://maven.apache.org/POM/4.0.0'}

    # 查找 <modules> 元素
    modules = root.find('ns:modules', namespaces=namespace)
    # 如果modules不存在
    if modules is None:
        modules = etree.Element("modules")

    # 判断module是否已存在，如果已存在，则从列表删除
    for module in modules.findall('ns:module', namespaces=namespace):
        if DEBUG_MODE:
            print(f"检查模块: {module.text}")
        if module.text in module_list:
            if DEBUG_MODE:
                print(f"已存在模块: {module.text}")
            add_list.remove(module.txt)
    # 新增列表中的元素
    for item in add_list:
        module = etree.SubElement(modules, "module")
        module.text = item

    # 写回文件，保留原始格式
    with open(pom_path, 'wb') as f:
        f.write(etree.tostring(tree, pretty_print=True, encoding='utf-8', xml_declaration=True))

    if DEBUG_MODE:
        print(f"POM 文件已更新: {pom_path}")

def main():
    command_parser = argparse.ArgumentParser(description="maven pom modifier script")
    # add args
    command_parser.add_argument("-d", "--dir", help="Modify a pom.xml in specify directory")
    command_parser.add_argument("-rm", "--removemodule", nargs="*", help="Remove <module>s from <modules>")
    command_parser.add_argument("-am", "--addmodule", nargs="*", help="Add <module>s from <modules>")

    args = command_parser.parse_args()
    print(args)

    work_dir = os.getcwd()
    if args.dir is not None:
        work_dir = os.path.abspath(args.dir)

    if args.removemodule != None:
        if DEBUG_MODE:
            print(f"remove_module参数数量为{len(args.removemodule)}")
        if len(args.removemodule) == 0:
            raise ValueError("At least one module must be specified")
        remove_module_list = args.removemodule
        remove_module(work_dir, remove_module_list)

    if args.addmodule != None:
        if DEBUG_MODE:
            print(f"add_module参数数量为{len(args.addmodule)}")
        if len(args.addmodule) == 0:
            raise ValueError("At least one module must be specified in add type")
        add_module_list = args.addmodule
        add_module(work_dir, add_module_list)

if __name__ == "__main__":
    main()